import json
import time
import re
import os
from openai import OpenAI
import datetime
from docx import Document

# ========================== 1. 配置区 ==========================
API_KEY = "BearerCubJNNWPCQjooxiEbLGB:IgSlCSEPDRuFXHJlXlMw"
URL = "https://spark-api-open.xf-yun.com/v2/chat/completions"
INPUT_DOCUMENT_PATH = r"C:/Users/cassi/Desktop/bisai/word-master/豫洛阳-兰郑长干线-CPY-0790-BFGDGS-ZZSYQFGS.docx"
OUTPUT_JSON_PATH = r"C:/Users/cassi/Desktop/bisai/word-master/temporal_logic_check.json"


# ========================== 2. 功能函数区 ==========================
def get_answer(message):
    """调用大模型生成负样本"""
    try:
        client = OpenAI(
            api_key="sk-ma0HIkN3beK0mtThtDzIfRLN5IzpL4agmB9RSbpHk53h7BF5",
            base_url="https://www.chataiapi.com/v1",
        )
        response = client.chat.completions.create(
            model="gemini-2.5-pro",
            messages=message,
            temperature=0.7,
        )
        content = response.choices[0].message.content or ""
        if content.strip():
            print("\n\n*******************模型生成内容如下********************\n")
            print(content)
        return content
    except Exception as e:
        print(f"API请求失败：{str(e)}")
        return ""


def extract_text_from_docx(docx_path):
    """提取Word文档中的纯文本内容"""
    try:
        doc = Document(docx_path)
        full_text = ""
        for para in doc.paragraphs:
            text = para.text.strip()
            if text:
                full_text += text + "\n"

        for table in doc.tables:
            for row in table.rows:
                row_text = []
                for cell in row.cells:
                    cell_text = cell.text.strip()
                    if cell_text:
                        row_text.append(cell_text)
                if row_text:
                    full_text += "  ".join(row_text) + "\n"

        # 🔧 自动修正跨行时间标识（关键修改点）
        # 例如 “识别时间\n2024年3月” → “识别时间：2024年3月”
        full_text = re.sub(r'(识别时间)\s*\n\s*(\d{4}年\d{1,2}月)', r'\1：\2', full_text)
        full_text = re.sub(r'(风险评价时间)\s*\n\s*(\d{4}年\d{1,2}月)', r'\1：\2', full_text)
        full_text = re.sub(r'(评价时间)\s*\n\s*(\d{4}年\d{1,2}月)', r'\1：\2', full_text)
        return full_text
    except Exception as e:
        print(f"Word文档文本提取失败：{str(e)}")
        return ""


def extract_date_from_text(date_text):
    """从文本中提取日期"""
    date_text = re.sub(r'[^\d年月日./\-\s]', '', date_text)
    date_patterns = [
        r'(\d{4})[年.-](\d{1,2})[月.-](\d{1,2})',
        r'(\d{4})[/.-](\d{1,2})[/.-](\d{1,2})',
        r'(\d{4})\s*年\s*(\d{1,2})\s*月',
        r'(\d{1,2})[月.-](\d{1,2})[日，]?\s*(\d{4})',
    ]
    for pattern in date_patterns:
        match = re.search(pattern, date_text)
        if match:
            try:
                groups = match.groups()
                if len(groups) == 3:
                    if len(groups[0]) == 4:
                        year, month, day = int(groups[0]), int(groups[1]), int(groups[2])
                    else:
                        month, day, year = int(groups[0]), int(groups[1]), int(groups[2])
                elif len(groups) == 2:
                    year, month, day = int(groups[0]), int(groups[1]), 1
                if 2000 <= year <= 2035 and 1 <= month <= 12 and 1 <= day <= 31:
                    return datetime.date(year, month, day)
            except:
                continue
    return None


def extract_time_information(text):
    """提取时间信息"""
    time_info = {
        'cover_dates': [],
        'identification_dates': [],
        'risk_assessment_dates': [],
        'other_dates': []
    }

    # 封面时间
    cover_patterns = [
        r'编制时间[：:\s]*([^\n]+)',
        r'制定时间[：:\s]*([^\n]+)',
        r'发布时间[：:\s]*([^\n]+)',
        r'版本时间[：:\s]*([^\n]+)'
    ]
    for pattern in cover_patterns:
        matches = re.finditer(pattern, text)
        for match in matches:
            date_text = match.group(1).strip()
            extracted_date = extract_date_from_text(date_text)
            if extracted_date:
                time_info['cover_dates'].append({
                    'raw_text': date_text,
                    'parsed_date': extracted_date,
                    'context': match.group(0),
                    'position': match.span()
                })

    # 识别时间
    identification_patterns = [
        r'识别时间[：:\s]*([^\n]*)',
        r'识别日期[：:\s]*([^\n]*)',
        r'完成识别.*?(\d{4}年\d{1,2}月)',
    ]
    for pattern in identification_patterns:
        matches = re.finditer(pattern, text)
        for match in matches:
            date_text = match.group(1).strip()
            extracted_date = extract_date_from_text(date_text)
            if extracted_date:
                time_info['identification_dates'].append({
                    'raw_text': date_text,
                    'parsed_date': extracted_date,
                    'context': match.group(0),
                    'position': match.span()
                })

    # 风险评价时间
    assessment_patterns = [
        r'风险评价时间[：:\s]*([^\n]*)',
        r'风险评价日期[：:\s]*([^\n]*)',
        r'评价时间[：:\s]*([^\n]*)',
        r'完成评价.*?(\d{4}年\d{1,2}月)',
    ]
    for pattern in assessment_patterns:
        matches = re.finditer(pattern, text)
        for match in matches:
            date_text = match.group(1).strip()
            extracted_date = extract_date_from_text(date_text)
            if extracted_date:
                time_info['risk_assessment_dates'].append({
                    'raw_text': date_text,
                    'parsed_date': extracted_date,
                    'context': match.group(0),
                    'position': match.span()
                })

    # ✅ 新增：只保留第一个风险评价时间，避免重复识别
    if len(time_info['risk_assessment_dates']) > 1:
        print(f"⚠️ 检测到多个风险评价时间（共{len(time_info['risk_assessment_dates'])}个），仅保留第一个。")
        time_info['risk_assessment_dates'] = [time_info['risk_assessment_dates'][0]]

    return time_info


def build_time_snippets(text, time_info, context_window=120):
    """根据识别到的时间位置信息，从原文中切片生成上下文片段。"""
    snippets = []
    for key in ['cover_dates', 'identification_dates', 'risk_assessment_dates']:
        for item in time_info.get(key, []):
            start, end = item.get('position', (0, 0))
            s = max(0, start - context_window)
            e = min(len(text), end + context_window)
            snippet = text[s:e].strip()
            if snippet and snippet not in snippets:
                snippets.append(snippet)
    # 若没有定位到具体位置，则回退使用前 N 字符
    if not snippets and text:
        snippets.append(text[: max(300, min(800, len(text)))] )
    return snippets


def format_year_month(d):
    """将日期格式化为YYYY年M月。"""
    return f"{d.year}年{d.month}月"


def build_canonical_time_line(time_info):
    """构造规范单行文本：封面编制时间：xx，高后果区识别时间：xx，风险评价时间：xx"""
    cover = time_info.get('cover_dates', [])
    ident = time_info.get('identification_dates', [])
    assess = time_info.get('risk_assessment_dates', [])

    cover_str = format_year_month(cover[0]['parsed_date']) if cover else 'xx'
    ident_str = format_year_month(ident[0]['parsed_date']) if ident else 'xx'
    assess_str = format_year_month(assess[0]['parsed_date']) if assess else 'xx'

    return f"封面编制时间：{cover_str}，高后果区识别时间：{ident_str}，风险评价时间：{assess_str}"


def get_temporal_negative_prompt(canonical_line, context_text, variant="order_conflict"):
    """构造“时间逻辑对抗样本”提示词：包含规范单行与原始上下文片段两部分。"""
    variant_to_desc = {
        # 让“风险评价时间”未晚于“高后果区识别时间”，违背时间逻辑
        'order_conflict': (
            "任务：生成‘时间先后顺序不合理’的对抗样本。\n"
            "要求：\n"
            "1. 原始文本为规范单行：‘封面编制时间：YYYY年M月，高后果区识别时间：YYYY年M月，风险评价时间：YYYY年M月’。保持结构与字段名不变。\n"
            "2. 尽量不改年份（YYYY），优先仅改月份/日，使‘风险评价时间’未晚于‘高后果区识别时间’（可同月或更早）。\n"
            "3. 不修改非时间文本与实体名称。\n"
            "原始文本（规范单行）：\n---\n{canonical}\n---\n\n原始上下文片段：\n---\n{context}\n---\n"
        ),
        # 年份不一致（仅修改年份），指定不一致的两类时间
        'year_inconsistency_cover_identification': (
            "任务：生成‘年份不一致（封面编制时间 vs 高后果区识别时间）’的对抗样本。\n"
            "要求：\n"
            "1. 原始文本为规范单行，保持结构与字段名不变。\n"
            "2. 仅修改年份（YYYY），使‘封面编制时间’与‘高后果区识别时间’年份不同（如封面2023，识别2024）。\n"
            "3. 月份尽量保持不变；不修改非时间文本。\n"
            "4. 输出仅为修改后的文本。\n"
            "原始文本（规范单行）：\n---\n{canonical}\n---\n\n原始上下文片段：\n---\n{context}\n---\n"
        ),
        'year_inconsistency_cover_assessment': (
            "任务：生成‘年份不一致（封面编制时间 vs 风险评价时间）’的对抗样本。\n"
            "要求：\n"
            "1. 原始文本为规范单行，保持结构与字段名不变。\n"
            "2. 仅修改年份（YYYY），使‘封面编制时间’与‘风险评价时间’年份不同。\n"
            "3. 月份尽量保持不变；不修改非时间文本。\n"
            "4. 输出仅为修改后的文本。\n"
            "原始文本（规范单行）：\n---\n{canonical}\n---\n\n原始上下文片段：\n---\n{context}\n---\n"
        ),
        'year_inconsistency_identification_assessment': (
            "任务：生成‘年份不一致（高后果区识别时间 vs 风险评价时间）’的对抗样本。\n"
            "要求：\n"
            "1. 原始文本为规范单行，保持结构与字段名不变。\n"
            "2. 仅修改年份（YYYY），使‘高后果区识别时间’与‘风险评价时间’年份不同。\n"
            "3. 月份尽量保持不变；不修改非时间文本。\n"
            "4. 输出仅为修改后的文本。\n"
            "原始文本（规范单行）：\n---\n{canonical}\n---\n\n原始上下文片段：\n---\n{context}\n---\n"
        ),
    }
    head = variant_to_desc.get(variant, variant_to_desc['order_conflict'])
    # 用具体原始文本填充模板
    prompt_filled = head.format(canonical=canonical_line or "", context=context_text or "")
    return f"{prompt_filled}\n输出：仅返回修改后的文本，无需额外解释。"


def check_temporal_logic(time_info):
    """检查时间逻辑一致性"""
    logic_results = {
        'year_consistency': True,
        'temporal_order_correct': True,
        'is_correct': True,
        'issues': [],
        'time_analysis': {}
    }

    all_dates = {
        'cover': [d['parsed_date'] for d in time_info['cover_dates']],
        'identification': [d['parsed_date'] for d in time_info['identification_dates']],
        'risk_assessment': [d['parsed_date'] for d in time_info['risk_assessment_dates']]
    }

    logic_results['time_analysis'] = {
        'cover_dates_count': len(all_dates['cover']),
        'identification_dates_count': len(all_dates['identification']),
        'risk_assessment_dates_count': len(all_dates['risk_assessment'])
    }

    # 年份一致性（细分三类两两不一致）
    cover_years = set(d.year for d in all_dates['cover'])
    identification_years = set(d.year for d in all_dates['identification'])
    risk_years = set(d.year for d in all_dates['risk_assessment'])

    # 任意两类均存在时进行对比
    if cover_years and identification_years and cover_years != identification_years:
        logic_results['year_consistency'] = False
        logic_results['is_correct'] = False
        logic_results['issues'].append(
            f"年份不一致（封面编制时间 vs 高后果区识别时间）：封面{sorted(cover_years)}，识别{sorted(identification_years)}"
        )
    if cover_years and risk_years and cover_years != risk_years:
        logic_results['year_consistency'] = False
        logic_results['is_correct'] = False
        logic_results['issues'].append(
            f"年份不一致（封面编制时间 vs 风险评价时间）：封面{sorted(cover_years)}，评价{sorted(risk_years)}"
        )
    if identification_years and risk_years and identification_years != risk_years:
        logic_results['year_consistency'] = False
        logic_results['is_correct'] = False
        logic_results['issues'].append(
            f"年份不一致（高后果区识别时间 vs 风险评价时间）：识别{sorted(identification_years)}，评价{sorted(risk_years)}"
        )

    # 时间先后逻辑
    if all_dates['identification'] and all_dates['risk_assessment']:
        if min(all_dates['risk_assessment']) <= max(all_dates['identification']):
            logic_results['temporal_order_correct'] = False
            logic_results['is_correct'] = False
            logic_results['issues'].append("风险评价时间未晚于高后果区识别时间")

    # 时间合理性
    current_date = datetime.date.today()
    for date_type, dates in all_dates.items():
        for d in dates:
            if d > current_date:
                logic_results['issues'].append(f"{date_type}时间({d})为未来时间")
                logic_results['is_correct'] = False

    # 时间缺失
    if not all_dates['cover']:
        logic_results['issues'].append("缺少封面编制时间")
        logic_results['is_correct'] = False
    if not all_dates['identification']:
        logic_results['issues'].append("缺少高后果区识别时间")
        logic_results['is_correct'] = False
    if not all_dates['risk_assessment']:
        logic_results['issues'].append("缺少风险评价时间")
        logic_results['is_correct'] = False

    return logic_results


# ========================== 主程序 ==========================
if __name__ == '__main__':
    file_name = os.path.basename(INPUT_DOCUMENT_PATH)
    print(f"准备处理文档: {file_name}")
    text = extract_text_from_docx(INPUT_DOCUMENT_PATH)
    print(f"成功提取文本，长度: {len(text)} 字符")

    time_info = extract_time_information(text)
    print(f"提取到时间信息：封面{len(time_info['cover_dates'])}个，识别{len(time_info['identification_dates'])}个，评价{len(time_info['risk_assessment_dates'])}个")

    logic = check_temporal_logic(time_info)
    print(f"检测结果: {logic['issues'] if logic['issues'] else '无问题'}")
    print("\n================ 时间信息明细 ================")
    for key, items in time_info.items():
        print(f"\n【{key}】共 {len(items)} 个")
        for i, item in enumerate(items, 1):
            print(f"  {i}. 原文片段：{item['context']}")
            print(f"     解析结果：{item['parsed_date']}")

    # 打印调试：年份与月份对比摘要
    def ym(d):
        return f"{d.year:04d}-{d.month:02d}"
    cover_dates = [ym(d) for d in [x['parsed_date'] for x in time_info['cover_dates']]]
    ident_dates = [ym(d) for d in [x['parsed_date'] for x in time_info['identification_dates']]]
    assess_dates = [ym(d) for d in [x['parsed_date'] for x in time_info['risk_assessment_dates']]]
    print("\n—— 年月对比（YYYY-MM）——")
    print(f"封面编制时间：{cover_dates if cover_dates else '无'}")
    print(f"高后果区识别时间：{ident_dates if ident_dates else '无'}")
    print(f"风险评价时间：{assess_dates if assess_dates else '无'}")

    # =============== 生成“时间逻辑”对抗样本（输出结构参考 image_text_consistency.py） ===============
    print("\n开始生成‘时间逻辑’类对抗样本……")
    # 使用规范单行 + 原始上下文片段作为“原始文本”；纠正提示仍以规范单行为准
    canonical_line = build_canonical_time_line(time_info)
    context_snippets = build_time_snippets(text, time_info)
    context_text = "\n---\n".join(context_snippets[:1]) if context_snippets else ""
    # 将规范单行与原始上下文分别传入提示词构造器
    snippets = [(canonical_line, context_text)]
    variants = [
        'order_conflict',
        'year_inconsistency_cover_identification',
        'year_inconsistency_cover_assessment',
        'year_inconsistency_identification_assessment',
    ]

    # 与 image_text_consistency.py 保持一致的输出结构：messages + images
    all_generated_data = []

    variant_to_cn = {
        'order_conflict': '先后顺序不合理',
        'year_inconsistency_cover_identification': '年份不一致（封面 vs 识别）',
        'year_inconsistency_cover_assessment': '年份不一致（封面 vs 评价）',
        'year_inconsistency_identification_assessment': '年份不一致（识别 vs 评价）',
    }

    for idx, (canonical_for_prompt, context_for_prompt) in enumerate(snippets):
        concat_len = len((canonical_for_prompt or "") + (context_for_prompt or ""))
        print(f"处理片段 {idx+1}/{len(snippets)}，长度：{concat_len}")
        source_label = f"{file_name} 片段{idx+1}"
        assistant_reference = canonical_line
        for var in variants:
            prompt = get_temporal_negative_prompt(canonical_for_prompt, context_for_prompt, variant=var)
            chat_message = [{"role": "user", "content": prompt}]
            negative_sample = get_answer(chat_message)
            if negative_sample.strip():
                issue_cn = variant_to_cn.get(var, var)
                # 构造更具体的不一致描述
                cover = time_info.get('cover_dates', [])
                ident = time_info.get('identification_dates', [])
                assess = time_info.get('risk_assessment_dates', [])
                cover_str = format_year_month(cover[0]['parsed_date']) if cover else 'xx'
                ident_str = format_year_month(ident[0]['parsed_date']) if ident else 'xx'
                assess_str = format_year_month(assess[0]['parsed_date']) if assess else 'xx'

                detail = issue_cn
                if var == 'year_inconsistency_cover_identification':
                    detail = f"年份不一致（封面编制时间和高后果区识别时间）"
                elif var == 'year_inconsistency_cover_assessment':
                    detail = f"年份不一致（封面编制时间和风险评价时间）"
                elif var == 'year_inconsistency_identification_assessment':
                    detail = f"年份不一致（高后果区识别时间和风险评价时间）"
                elif var == 'order_conflict':
                    detail = f"时间先后顺序不合理（风险评价时间未晚于高后果区识别时间）"

                all_generated_data.append({
                    "messages": [
                        {
                            "content": f"{negative_sample.strip()}",
                            "role": "user"
                        },
                        {
                            "content": f"检测到时间逻辑问题：{detail}。正确的描述应该是：\n{assistant_reference}",
                            "role": "assistant"
                        }
                    ],
                    "images": []
                })
            else:
                print(f"  → 变体 {var} 生成失败，跳过。")
            time.sleep(0.6)

    # 持久化输出（结构同 image_text_consistency.py）
    if all_generated_data:
        try:
            out_dir = os.path.dirname(OUTPUT_JSON_PATH)
            if out_dir and not os.path.exists(out_dir):
                os.makedirs(out_dir)
            with open(OUTPUT_JSON_PATH, "w", encoding="utf-8") as f:
                json.dump(all_generated_data, f, ensure_ascii=False, indent=4)
            print("\n" + "="*80)
            print(f"✅ 处理完成！数据保存至：{OUTPUT_JSON_PATH}")
            print(f"共生成 {len(all_generated_data)} 个对话样本（每个包含负样本user + 正样本assistant）")
            print("="*80 + "\n")
        except Exception as e:
            print(f"保存对抗样本失败：{e}")
    else:
        print("\n⚠️ 处理结束，无生成数据。请检查文本片段或API配置。")
